#!/bin/bash
# 定义 Docker 仓库的基础 URL
# DOCKER_REPO_BASE="https://download.docker.com/linux"
DOCKER_REPO_BASE="https://mirrors.ustc.edu.cn/docker-ce/linux"
# 检测处理器架构
detect_arch() {
    ARCH=$(uname -m)
    case $ARCH in
        x86_64)
            ARCH="amd64"
            ;;
        aarch64)
            ARCH="arm64"
            ;;
        armv7l)
            ARCH="armhf"
            ;;
        *)
            echo "不支持的架构: $ARCH"
            exit 1
            ;;
    esac
}

# 获取 Linux 发行版信息
detect_os() {
    DISTRO=$(lsb_release -is | tr '[:upper:]' '[:lower:]')
    CODENAME=$(lsb_release -cs)
}

install_docker() {
    # 安装依赖
    sudo apt-get update -y
    sudo apt install -y software-properties-common curl apt-transport-https ca-certificates gnupg lsb-release

    # 定义 Docker 仓库的 URL
    DOCKER_REPO_URL="${DOCKER_REPO_BASE}/${DISTRO} ${CODENAME}"
    # 检测架构并设置仓库 URL（必要时可以扩展支持更多架构）
    case $ARCH in
        amd64|x86_64)
            DOCKER_REPO_URL="${DOCKER_REPO_URL} stable"
            ;;
        aarch64|arm64)
            DOCKER_REPO_URL="${DOCKER_REPO_URL} testing" # 假设 ARM 架构使用测试仓库
            ;;
        *)
            echo "Unsupported architecture: $ARCH"
            exit 1
            ;;
    esac
    # 打印将要使用的 Docker 仓库 URL
    echo "Using Docker repository: $DOCKER_REPO_URL"
    GPG_PATH=/usr/share/keyrings/docker-archive-keyring.gpg
    # 添加 Docker 的 GPG 密钥
    curl -fsSL "${DOCKER_REPO_BASE}/${DISTRO}/gpg" | sudo gpg --dearmor -o ${GPG_PATH}
    # 向 sources.list.d 添加 Docker 仓库
    echo "deb [arch=$ARCH signed-by=${GPG_PATH}] $DOCKER_REPO_URL" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
    # 更新 apt 包索引
    sudo apt update -y
    sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
    newgrp docker
    sudo usermod -aG docker $USER
}

install_nvidia_docker() {
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
    && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container.list
    sudo apt update -y
    sudo apt install -y nvidia-container-toolkit
}

detect_os
detect_arch

case $DISTRO in
    "ubuntu"|"debian")
        install_docker
        
        if [ "${1}" = "nvidia" ];then
            install_nvidia_docker
        fi
        echo "Docker repository 已添加"
        ;;
    *)
        echo "不支持的操作系统: $DISTRO"
        exit 1
        ;;
esac